package com.shujia.spark.core

import org.apache.spark.{Partitioner, SparkConf, SparkContext}
import org.apache.spark.rdd.RDD

object Demo21Partitions {
  /**
   * 为什么要讨论RDD中分区数？
   * Task是作用在每个分区上的，改变分区的数量可以
   * 间接的影响Task的数量，进而影响整个的任务的运行效率
   *
   */
  def main(args: Array[String]): Unit = {
    val conf: SparkConf = new SparkConf()
    conf.setMaster("local[4]")
    conf.setAppName("Demo21Partitions")
    conf.set("spark.default.parallelism", "3")

    val sc: SparkContext = new SparkContext(conf)

    // 默认等于切片的数量
    /**
     * 当Spark任务的并行度超过2时，要求RDD的最小分区数为2
     * 如果只有一个文件比较小，对应只有一个切分
     * 通常会对应一个分区，在并行度超过2的情况下，会自动变成2个分区来满足最小分区数的要求
     *
     * 如果不止一个文件，也就意味着切片的数量会大于等于2
     * 那这个时候就以切片的数量为准来决定第一个RDD的分区数
     *
     * 或者从并行度的角度看：
     * 如果并行度=1，则RDD的分区的数量等于切片的数量
     * 如果并行度>1，则RDD的分区数最少为2个
     */
    val stuRDD: RDD[String] = sc.textFile("spark/data/stu/students.txt")
    println(s"stuRDD的分区数为:${stuRDD.getNumPartitions}")

    // 统计班级人数
    val clazzRDD: RDD[(String, Int)] = stuRDD.map(line => (line.split(",")(4), 1))
    // 没有shuffle操作，子RDD的分区数等于父RDD的分区数
    println(s"clazzRDD的分区数为:${clazzRDD.getNumPartitions}")

    val clazzCntRDD01: RDD[(String, Int)] = clazzRDD.reduceByKey(_ + _)
    val clazzCntRDD02: RDD[(String, Int)] = clazzRDD.reduceByKey(_ + _, numPartitions = 4)
    val clazzCntRDD03: RDD[(String, Int)] = clazzRDD.reduceByKey(new MyPartitioner(12), _ + _)

    /**
     * 有Shuffle的操作:
     * 1、默认情况下还是等于父RDD的分区数
     * 2、可以在算子中通过numPartitions参数或者是指定一个Partitioner来决定子RDD的分区数
     * 3、如果指定了spark.default.parallelism参数，则默认子RDD分区数由该参数决定
     * 优先级：
     * 手动指定numPartitions或者自定分区 >> spark.default.parallelism参数 >> 默认等于父RDD的分区数
     *
     */
    println(s"clazzCntRDD01的分区数为:${clazzCntRDD01.getNumPartitions}")
    println(s"clazzCntRDD02的分区数为:${clazzCntRDD02.getNumPartitions}")
    println(s"clazzCntRDD03的分区数为:${clazzCntRDD03.getNumPartitions}")

    // 将结果保存到文件中
    //    clazzCntRDD01.saveAsTextFile("spark/data/clazzCnt/01/")
    //    clazzCntRDD02.saveAsTextFile("spark/data/clazzCnt/02/")
    //    clazzCntRDD03.saveAsTextFile("spark/data/clazzCnt/03/")

    // 可以通过repartition算子进行重分区
    /**
     * repartition实际上就是coalesce将shuffle参数设为true
     * 表示通过shuffle改变分区数
     */
    val repartitionRDD01: RDD[(String, Int)] = clazzCntRDD02.repartition(2)
    val repartitionRDD02: RDD[(String, Int)] = clazzCntRDD02.repartition(8)
    println(s"repartitionRDD01的分区数为:${repartitionRDD01.getNumPartitions}")
    println(s"repartitionRDD02的分区数为:${repartitionRDD02.getNumPartitions}")

    // 还可以通过coalesce改变分区的数量
    /**
     * coalesce默认shuffle为false，所以只能用于减少分区
     * 如果需要增加分区数则需要使用Shuffle完成
     */
    val coalesceRDD01: RDD[(String, Int)] = clazzCntRDD02.coalesce(2)
    val coalesceRDD02: RDD[(String, Int)] = clazzCntRDD02.coalesce(8)
    println(s"coalesceRDD01的分区数为:${coalesceRDD01.getNumPartitions}")
    println(s"coalesceRDD02的分区数为:${coalesceRDD02.getNumPartitions}")
  }

}

// 自定义分区器
class MyPartitioner(num: Int) extends Partitioner {
  // 指定分区数量
  override def numPartitions: Int = num

  override def getPartition(key: Any): Int = {
    val clazz: String = key.toString
    clazz match {
      case "文科一班" => 0
      case "文科二班" => 1
      case "文科三班" => 2
      case "文科四班" => 3
      case "文科五班" => 4
      case "文科六班" => 5
      case "理科一班" => 6
      case "理科二班" => 7
      case "理科三班" => 8
      case "理科四班" => 9
      case "理科五班" => 10
      case "理科六班" => 11
    }
  }
}
